package main

import (
	"encoding/json"
	"io"
	"os"
)

// 只保留我们关心的字段，其余原样透传
type notebook struct {
	Cells []struct {
		CellType       string        `json:"cell_type"`
		ExecutionCount interface{}   `json:"execution_count,omitempty"`
		Outputs        []interface{} `json:"outputs,omitempty"`
	} `json:"cells"`
	// 其它顶层字段（metadata、nbformat 等）原样保留
	Metadata      interface{} `json:"metadata"`
	Nbformat      int         `json:"nbformat"`
	NbformatMinor int         `json:"nbformat_minor"`
}

func main() {
	if err := strip(os.Stdin, os.Stdout); err != nil {
		os.Stderr.WriteString("nbstrip: " + err.Error() + "\n")
		os.Exit(1)
	}
}

func strip(r io.Reader, w io.Writer) error {
	var nb notebook
	if err := json.NewDecoder(r).Decode(&nb); err != nil {
		return err
	}

	// 清空 code cell 的输出
	for i := range nb.Cells {
		c := &nb.Cells[i]
		if c.CellType == "code" {
			c.Outputs = []interface{}{}
			c.ExecutionCount = nil
		}
	}

	// 重新编码，保持缩进与原格式一致
	enc := json.NewEncoder(w)
	enc.SetIndent("", " ")
	return enc.Encode(&nb)
}

/*
go build -ldflags="-s -w" -trimpath -o nbstrip nbstrip.go
# 可选再压
upx --best nbstrip   # 可再减 40 %
*/
